
#ifndef AFTERMATH_NUMERICS_DSSE_MATRIX_HPP_INCLUDED
#define AFTERMATH_NUMERICS_DSSE_MATRIX_HPP_INCLUDED

#include "../aligned_memory.hpp"
#include "dsse_math.hpp"
#include "matrix.hpp"

#include <emmintrin.h>
#include <memory.h>
#include <mmintrin.h>
#include <smmintrin.h>
#include <utility>

namespace aftermath
{
    namespace numerics
    {
        template <int M, int N>
        struct matrix<M, N, double, __m128d>
        {
            template <int K, int L, typename T, typename S>
            friend struct matrix;

            typedef double data_type;
            typedef __m128d sse_type;

            static const int height = M;
            static const int width = N;
            
            static const bool is_square = (M == N);
            static const int memory_wasted = (N % 2) * sizeof(data_type);
        
            static const int required_align = 16;
            static const int stride = (N + 1) / 2;
            static const int required_memory = M * stride * sizeof(sse_type);
            static const int sufficient_memory = sufficient_memory_size<required_memory, required_align>::value;
            
            matrix()
                : m_memory(nullptr), m_data(nullptr)
            {
                //this->m_memory = new char[sufficient_memory];
                this->m_memory = malloc(sufficient_memory);
                this->m_data = aligned_memory<sse_type, required_align>::address(this->m_memory);
                /** Clean border. */
                if (N % 2 != 0) for (int i = 0; i < M; i++) this->get_sse_element(i, N / 2) = _mm_setzero_pd();
            }

            matrix(const matrix& other)
                : m_memory(nullptr), m_data(nullptr)
            {
                //this->m_memory = new char[sufficient_memory];
                this->m_memory = malloc(sufficient_memory);
                this->m_data = aligned_memory<sse_type, required_align>::address(this->m_memory);
                std::memcpy(this->m_data, other.m_data, required_memory); 
            }

            matrix(matrix&& other)
                : m_memory(nullptr), m_data(nullptr)
            {
                *this = std::move(other);
            }

            matrix& operator =(const matrix& other) 
            {
                if (this != &other) std::memcpy(this->m_data, other.m_data, required_memory);
                return *this;
            }

            matrix& operator =(matrix&& other)
            {
                if (this != &other)
                {
                    free(this->m_memory);
                    this->m_data = other.m_data;
                    this->m_memory = other.m_memory;
                    other.m_data = nullptr;
                    other.m_memory = nullptr;
                }
                return *this;
            }

            ~matrix()
            {
                //delete[] this->m_memory;
                free(this->m_memory);
                this->m_memory = nullptr;
                this->m_data = nullptr;
            }

            bool is_neat() const
            {
                if (N % 2 == 0) return true; 
                for (int i = 0; i < M; i++) if (this->get_sse_element(i, N / 2).m128d_f64[1] != 0.0) return false;
                return true;
            }
            
            void fill(data_type value)
            {
                __m128d v = _mm_loaddup_pd(&value);
                if (N % 2 == 0) for (int i = 0; i < M; i++) for (int j = 0; j < N / 2; j++) this->get_sse_element(i, j) = v;
                else
                {
                    __m128d w = _mm_load_sd(&value);
                    for (int i = 0; i < M; i++) 
                    {
                        for (int j = 0; j < N / 2; j++) this->get_sse_element(i, j) = v;
                        this->get_sse_element(i, N / 2) = w;
                    }
                }
            }
            
            void fill_row(int row, data_type value) 
            {
                __m128d v = _mm_loaddup_pd(&value);
                for (int j = 0; j < N / 2; j++) this->get_sse_element(row, j) = v;
                if (N % 2 == 1) this->get_sse_element(row, N / 2) = _mm_load_sd(&value);
            }

            void clear() { for (int i = 0; i < M; i++) for (int j = 0; j < (N + 1) / 2; j++) this->get_sse_element(i, j) = _mm_setzero_pd(); }

            void clear_row(int row) { for (int j = 0; j < (N + 1) / 2; j++) this->get_sse_element(row, j) = _mm_setzero_pd(); }
        
            /** Elementwise operations are slow! */
            data_type get_element(int row, int column) const { return this->get_sse_element(row, column / 2).m128d_f64[column % 2]; }

            template <int row, int column>
            void set_element(data_type value)
            {
                //__declspec(align(16)) double v = value;
                if (column % 2 == 0) this->get_sse_element(row, column / 2) = _mm_unpackhi_pd(_mm_loaddup_pd(&value), this->get_sse_element(row, column / 2));
                if (column % 2 == 1) this->get_sse_element(row, column / 2) = _mm_unpacklo_pd(this->get_sse_element(row, column / 2), _mm_loaddup_pd(&value));
            }

            void set_element(int row, int column, data_type value)
            {
                //__declspec(align(16)) double v = value;
                if (column % 2 == 0) this->get_sse_element(row, column / 2) = _mm_unpackhi_pd(_mm_loaddup_pd(&value), this->get_sse_element(row, column / 2));
                if (column % 2 == 1) this->get_sse_element(row, column / 2) = _mm_unpacklo_pd(this->get_sse_element(row, column / 2), _mm_loaddup_pd(&value));
            }
            
            template <int row, int column>
            void clear_element()
            {
                if (column % 2 == 0) this->get_sse_element(row, column / 2) = _mm_unpackhi_pd(_mm_setzero_pd(), this->get_sse_element(row, column / 2));
                if (column % 2 == 1) this->get_sse_element(row, column / 2) = _mm_unpacklo_pd(this->get_sse_element(row, column / 2), _mm_setzero_pd());
            }
            
            void clear_element(int row, int column)
            {
                if (column % 2 == 0) this->get_sse_element(row, column / 2) = _mm_unpackhi_pd(_mm_setzero_pd(), this->get_sse_element(row, column / 2));
                if (column % 2 == 1) this->get_sse_element(row, column / 2) = _mm_unpacklo_pd(this->get_sse_element(row, column / 2), _mm_setzero_pd());
            }
            
            void transpose_self();
            matrix<N, M, double, __m128d> transpose() const;

            template <int K>
            matrix<M, K, double, __m128d> multiply(const matrix<N, K, double, __m128d>& other) const;

            template <int K>
            matrix<M, K, double, __m128d> fast_multiply(const matrix<N, K, double, __m128d>& other) const;
            
            matrix<M, N, double, __m128d> operator +(const matrix<M, N, double, __m128d>& other)
            {
                matrix<M, N, double, __m128d> result;
                for (int i = 0; i < M; i++) for (int j = 0; j < (N + 1) / 2; j++) result.get_sse_element(i, j) = this->get_sse_element(i, j) + other.get_sse_element(i, j);
                return result;
            }
            
            matrix<M, N, double, __m128d> operator -(const matrix<M, N, double, __m128d>& other)
            {
                matrix<M, N, double, __m128d> result;
                for (int i = 0; i < M; i++) for (int j = 0; j < (N + 1) / 2; j++) result.get_sse_element(i, j) = this->get_sse_element(i, j) - other.get_sse_element(i, j);
                return result;
            }

            matrix<M, N, double, __m128d> operator *(const matrix<M, N, double, __m128d>& other)
            {
                matrix<M, N, double, __m128d> result;
                for (int i = 0; i < M; i++) for (int j = 0; j < (N + 1) / 2; j++) result.get_sse_element(i, j) = this->get_sse_element(i, j) * other.get_sse_element(i, j);
                return result;
            }

            matrix<M, N, double, __m128d> operator /(const matrix<M, N, double, __m128d>& other)
            {
                matrix<M, N, double, __m128d> result;
                for (int i = 0; i < M; i++) for (int j = 0; j < (N + 1) / 2; j++) result.get_sse_element(i, j) = this->get_sse_element(i, j) / other.get_sse_element(i, j);
                return result;
            }

            matrix<M, N, double, __m128d>& operator +=(const matrix<M, N, double, __m128d>& other)
            {
                for (int i = 0; i < M; i++) for (int j = 0; j < (N + 1) / 2; j++) this->get_sse_element(i, j) += other.get_sse_element(i, j);
                return *this;
            }

            matrix<M, N, double, __m128d>& operator -=(const matrix<M, N, double, __m128d>& other)
            {
                for (int i = 0; i < M; i++) for (int j = 0; j < (N + 1) / 2; j++) this->get_sse_element(i, j) -= other.get_sse_element(i, j);
                return *this;
            }

            matrix<M, N, double, __m128d>& operator *=(const matrix<M, N, double, __m128d>& other)
            {
                for (int i = 0; i < M; i++) for (int j = 0; j < (N + 1) / 2; j++) this->get_sse_element(i, j) *= other.get_sse_element(i, j);
                return *this;
            }

            matrix<M, N, double, __m128d>& operator /=(const matrix<M, N, double, __m128d>& other)
            {
                for (int i = 0; i < M; i++) for (int j = 0; j < (N + 1) / 2; j++) this->get_sse_element(i, j) /= other.get_sse_element(i, j);
                return *this;
            }

        private:
            // Row-major
            void* m_memory;
            __m128d* m_data;

            int get_sse_index(int i, int j) const { return i * stride + j; }

            __m128d& get_sse_element(int i, int j) const { return this->m_data[i * stride + j]; }
        };

        template <int M, int N>
        void matrix<M, N, double, __m128d>::transpose_self()
        {
            static_assert(M == N, "transpose_self only works on square matrices.");

            __m128d tmp1, tmp0;
            __m128d zero = _mm_setzero_pd();
            const int nBlocks = N / 2; // Number of full 2-by-2 blocks.

            // Core of the matrix: nBlocks-by-nBlocks block matrix with 2-by-2 blocks
            // Along the main diagonal do 2-by-2 transpose
            // For the off-diagonal blocks do coupled 2-by-2 transpose
            // 0-------------------------
            // |  1----------------------
            // |  |  2-------------------
            // |  |  |  3----------------
            // |  |  |  |  ...           
            for (int i = 0; i < nBlocks; i++)
            {
                // Main diagonal: 2-by-2 transpose
                bi_row_transposeSelf<2>(
                    this->get_sse_element(2 * i, i), this->get_sse_element(2 * i + 1, i),
                    tmp0, zero);
            
                for (int j = i + 1; j < nBlocks; j++)
                {
                    // Off-diagonal: 2-by-2 coupled transpose
                    bi_coupled_row_transpose<2>(
                        this->get_sse_element(2 * i, j), this->get_sse_element(2 * i + 1, j),
                        this->get_sse_element(2 * j, i), this->get_sse_element(2 * j + 1, i),
                        tmp0, tmp1, zero);
                }

                // Margins
                if (N % 2 != 0)
                    bi_coupled_row_transpose<N % 2>(
                        this->get_sse_element(2 * i, nBlocks), this->get_sse_element(2 * i + 1, nBlocks),
                        this->get_sse_element(2 * nBlocks, i), zero,
                        tmp0, tmp1, zero);
            }
        }

        template <int M, int N>
        matrix<N, M, double, __m128d> matrix<M, N, double, __m128d>::transpose() const
        {
            __m128d zero = _mm_setzero_pd();
            matrix<N, M, double, __m128d> result;
            const int nRowBlocks = M / 2; // Number of full 2-by-2 blocks.
            const int nColumnBlocks = N / 2; // Number of full 2-by-2 blocks.

            // Core matrix: nRowBlocks-by-nColumnBlocks block matrix with 2-by-2 blocks
            for (int i = 0; i < nRowBlocks; i++)
            {
                for (int j = 0; j < nColumnBlocks; j++)
                {
                    bi_row_transpose<2, 2>(
                        this->get_sse_element(2 * i, j), this->get_sse_element(2 * i + 1, j),
                        result.get_sse_element(2 * j, i), result.get_sse_element(2 * j + 1, i),
                        zero);
                }
            
                // Last column
                if (N % 2 != 0)
                    bi_row_transpose<2, N % 2>(
                        this->get_sse_element(2 * i, nColumnBlocks), this->get_sse_element(2 * i + 1, nColumnBlocks),
                        result.get_sse_element(2 * nColumnBlocks, i), zero, 
                        zero);
            }
        
            // Last row
            if (M % 2 != 0)
            {
                for (int j = 0; j < nColumnBlocks; j++)
                {
                    bi_row_transpose<M % 2, 2>(
                        this->get_sse_element(2 * nRowBlocks, j), zero,
                        result.get_sse_element(2 * j, nRowBlocks), result.get_sse_element(2 * j + 1, nRowBlocks),
                        zero);
                }

                // Main diagonal appendix
                if (N % 2 != 0)
                    bi_row_transpose<M % 2, N % 2>(
                        this->get_sse_element(2 * nRowBlocks, nColumnBlocks), zero,
                        result.get_sse_element(2 * nColumnBlocks, nRowBlocks), zero, 
                        zero);
            }

            return result;
        }

        template <int M, int N>
        template <int K>
        matrix<M, K, double, __m128d> matrix<M, N, double, __m128d>::multiply(const matrix<N, K, double, __m128d>& other) const
        {
            __m128d tmp0, tmp1;
            __m128d col0, col1;
            __m128d zero = _mm_setzero_pd();
            matrix<M, K, double, __m128d> result;
            const int nRowBlocks = M / 2; // Number of full 2-by-2 blocks in resulting matrix.
            const int nColumnBlocks = K / 2; // Number of full 2-by-2 blocks in resulting matrix.
            const int nAuxBlocks = N / 2; // Number of full 2-by-2 blocks in auxiliary matrix.

            // Core matrix: nRowBlocks-by-nColumnBlocks block matrix with 2-by-2 blocks.
            // Since we need to transpose the 'other' matrix, we will be moving along colums
            // so as to minimize the number of transpositions.
            for (int j = 0; j < nColumnBlocks; j++)
            {
                // First (l = 0) we have to initialize result matrix and then cover the remaining blocks.
                if (nAuxBlocks > 0) // Check if nAuxBlocks > 0, i.e. N > 1
                {
                    bi_row_transpose<2, 2>(
                        other.get_sse_element(0, j), other.get_sse_element(1, j),
                        col0, col1, 
                        zero);

                    for (int i = 0; i < nRowBlocks; i++)
                    {
                        bi_multiply<2, 2, false>(
                                this->get_sse_element(2 * i, 0), this->get_sse_element(2 * i + 1, 0),
                                col0, col1,
                                result.get_sse_element(2 * i, j), result.get_sse_element(2 * i + 1, j),
                                tmp0, tmp1, zero);
                    }

                    // Last row
                    if (M % 2 != 0)
                        bi_multiply<M % 4, 2, false>(
                                this->get_sse_element(2 * nRowBlocks, 0), zero, 
                                col0, col1,
                                result.get_sse_element(2 * nRowBlocks, j), zero,
                                tmp0, tmp1, zero);
                }

                for (int l = 1; l < nAuxBlocks; l++)
                {
                    bi_row_transpose<2, 2>(
                        other.get_sse_element(2 * l, j), other.get_sse_element(2 * l + 1, j),
                        col0, col1, 
                        zero);

                    for (int i = 0; i < nRowBlocks; i++)
                    {
                        bi_multiply<2, 2, true>(
                                this->get_sse_element(2 * i, l), this->get_sse_element(2 * i + 1, l),
                                col0, col1,
                                result.get_sse_element(2 * i, j), result.get_sse_element(2 * i + 1, j),
                                tmp0, tmp1, zero);
                    }

                    // Last row
                    if (M % 2 != 0)
                        bi_multiply<M % 2, 2, true>(
                                this->get_sse_element(2 * nRowBlocks, l), zero, 
                                col0, col1,
                                result.get_sse_element(2 * nRowBlocks, j), zero,
                                tmp0, tmp1, zero);
                }

                // Last auxiliary row
                if (N % 2 != 0)
                {
                    bi_row_transpose<N % 2, 2>(
                        other.get_sse_element(2 * nAuxBlocks, j), zero,
                        col0, col1, 
                        zero);

                    for (int i = 0; i < nRowBlocks; i++)
                    {
                        bi_multiply<2, 2, nAuxBlocks != 0>(
                                this->get_sse_element(2 * i, nAuxBlocks), this->get_sse_element(2 * i + 1, nAuxBlocks),
                                col0, col1,
                                result.get_sse_element(2 * i, j), result.get_sse_element(2 * i + 1, j),
                                tmp0, tmp1, zero);
                    }

                    // Last row
                    if (M % 2 != 0)
                        bi_multiply<M % 2, 2, nAuxBlocks != 0>(
                                this->get_sse_element(2 * nRowBlocks, nAuxBlocks), zero, 
                                col0, col1,
                                result.get_sse_element(2 * nRowBlocks, j), zero,
                                tmp0, tmp1, zero);
                }
            }
        
            // Last column
            if (K % 2 != 0)
            {
                // First (l = 0) we have to initialize result matrix and then cover the remaining blocks.
                if (nAuxBlocks > 0) // Check if nAuxBlocks > 0, i.e. N > 1
                {
                    bi_row_transpose<2, K % 2>(
                        other.get_sse_element(0, nColumnBlocks), other.get_sse_element(1, nColumnBlocks),
                        col0, col1, 
                        zero);

                    for (int i = 0; i < nRowBlocks; i++)
                    {
                        bi_multiply<2, K % 2, false>(
                                this->get_sse_element(2 * i, 0), this->get_sse_element(2 * i + 1, 0),
                                col0, col1,
                                result.get_sse_element(2 * i, nColumnBlocks), result.get_sse_element(2 * i + 1, nColumnBlocks),
                                tmp0, tmp1, zero);
                    }

                    // Last row
                    if (M % 2 != 0)
                        bi_multiply<M % 2, K % 2, false>(
                                this->get_sse_element(2 * nRowBlocks, 0), zero, 
                                col0, col1,
                                result.get_sse_element(2 * nRowBlocks, nColumnBlocks), zero,
                                tmp0, tmp1, zero);
                }

                for (int l = 1; l < nAuxBlocks; l++)
                {
                    bi_row_transpose<2, K % 2>(
                        other.get_sse_element(2 * l, nColumnBlocks), other.get_sse_element(2 * l + 1, nColumnBlocks),
                        col0, col1, 
                        zero);

                    for (int i = 0; i < nRowBlocks; i++)
                    {
                        bi_multiply<2, K % 2, true>(
                                this->get_sse_element(2 * i, l), this->get_sse_element(2 * i + 1, l),
                                col0, col1,
                                result.get_sse_element(2 * i, nColumnBlocks), result.get_sse_element(2 * i + 1, nColumnBlocks),
                                tmp0, tmp1, zero);
                    }

                    // Last row
                    if (M % 2 != 0)
                        bi_multiply<M % 2, K % 2, true>(
                                this->get_sse_element(2 * nRowBlocks, l), zero, 
                                col0, col1,
                                result.get_sse_element(2 * nRowBlocks, nColumnBlocks), zero,
                                tmp0, tmp1, zero);
                }

                // Last auxiliary row
                if (N % 2 != 0)
                {
                    bi_row_transpose<N % 2, K % 2>(
                        other.get_sse_element(2 * nAuxBlocks, nColumnBlocks), zero,
                        col0, col1, 
                        zero);

                    for (int i = 0; i < nRowBlocks; i++)
                    {
                        bi_multiply<2, K % 2, true>(
                                this->get_sse_element(2 * i, nAuxBlocks), this->get_sse_element(2 * i + 1, nAuxBlocks),
                                col0, col1,
                                result.get_sse_element(2 * i, nColumnBlocks), result.get_sse_element(2 * i + 1, nColumnBlocks),
                                tmp0, tmp1, zero);
                    }

                    // Last row
                    if (M % 2 != 0)
                        bi_multiply<M % 2, K % 2, true>(
                                this->get_sse_element(2 * nRowBlocks, nAuxBlocks), zero, 
                                col0, col1,
                                result.get_sse_element(2 * nRowBlocks, nColumnBlocks), zero,
                                tmp0, tmp1, zero);
                }
            }

            return result;
        }
    }
}

#endif // AFTERMATH_NUMERICS_DSSE_MATRIX_HPP_INCLUDED
